import json, os, subprocess, time, stat, platform, importlib
import numpy as np
from Utils.network import TcpClientP2PWithCompress, find_free_port_no_repeat, get_host_ip
from Utils.config_args import ChainVar
from config import GlobalConfig
from Common.base_env import BaseEnv
from .action_set_encoder import binary_friendly, dictionary_n_actions
from .agent import Agent
import config


# please register this into Environment/route_and_environment.py

# def make_uhmap_env(env_id, rank):
#     if ScenarioConfig.SubTaskSelection == 'UhmapEnv':
#         return CssimEnv(rank)
#     else:
#         ST = ScenarioConfig.SubTaskSelection
#         assert os.path.exists(f'./Environment/cssim/{ST}.py'), "Unknown subtask!"
#         ST_CLASS = getattr(importlib.import_module(f'.cssim.{ST}', package='Environment'), ST)
#         return ST_CLASS(rank)
#
# def make_game_env(env_id, rank):
#     if ScenarioConfig.SubTaskSelection == 'UhmapEnv':
#         return CssimEnv(rank)
#     else:
#         ST = ScenarioConfig.SubTaskSelection
#         package_name, class_name = ST.split('.')
#         assert os.path.exists(f'./{GlobalConfig.env_path.replace(".","/")}/{ST}.py'), "Unknown subtask!"
#         ST_CLASS = getattr(importlib.import_module(f'{GlobalConfig.env_path}.{ST}', package=package_name), class_name)
#         return ST_CLASS(rank)
#
# def get_subtask_conf(subtask):
#     ST = subtask
#     assert os.path.exists(f'./Environment/cssim/{ST}Conf.py'), "Configuration not found!"
#     ST_CONF_CLASS = getattr(importlib.import_module(f'.cssim.{ST}Conf', package='Environment'), 'SubTaskConfig')
#     return ST_CONF_CLASS


def make_uhmap_env(env_id, rank, algorithms):
    ST = ScenarioConfig.SubTaskSelection
    assert '->' in ScenarioConfig.SubTaskSelection  ####
    # assert os.path.exists(f'./Environment/cssim/{ST}.py'), "Unknown subtask!"
    # ST_CLASS = getattr(importlib.import_module(f'.cssim.{ST}', package='Environment'), ST)

    file_, class_ = ST.split('->')
    env_path = file_[:-3].replace(".", "/") + '.py'
    assert os.path.exists(env_path), "Unknown Agent Environment!"
    if '.py' in file_:
        if file_.endswith('.py'): file_ = file_[:-3]
    ST_CLASS = getattr(importlib.import_module(file_), class_)

    return ST_CLASS(rank,algorithms)

def get_subtask_conf(subtask):
    ST = subtask
    assert '->' in ST  ####
    file_, class_ = ST.split('->')
    env_path = file_[:-3].replace(".", "/") + '.py'
    assert os.path.exists(env_path), "Unknown Agent Environment!"
    # env_conf_path = f'Algorithm/common/' + class_ + 'Conf.py'  # 2025/07
    env_conf_path = f'Common/' + class_ + 'Conf.py'
    assert os.path.exists(env_conf_path), "Environment Config not found!"
    env_conf_path = env_conf_path.replace("/", ".")[:-3]
    ST_CONF_CLASS = getattr(importlib.import_module(env_conf_path), 'SubTaskConfig')

    return ST_CONF_CLASS


def usual_id_arrangment(N_AGENT_EACH_TEAM):
    """
        e.g., 
        input [5, 3]
        output [range(0,5), range(5,8)]
    """
    AGENT_ID_EACH_TEAM = []
    p = 0
    for team_agent_num in N_AGENT_EACH_TEAM:
        AGENT_ID_EACH_TEAM.append(range(p, p + team_agent_num))
        p += team_agent_num
    return AGENT_ID_EACH_TEAM


# please register this ScenarioConfig into Environment/route_and_environment.py
class ScenarioConfig(object):
    '''
        ScenarioConfig: This config class will be 'injected' with new settings from JSONC.
        (E.g., override configs with ```python main.py --cfg example.jsonc```)
        (As the name indicated, ChainVars will change WITH vars it 'chained_with' during config injection)
        (please see Utils.config_args to find out how this advanced trick works out.)
    '''
    # <Part 1> Needed by the hmp core #
    # N_AGENT_EACH_TEAM = [25, 35]
    N_AGENT_EACH_TEAM = [5, 7]
    AGENT_ID_EACH_TEAM = usual_id_arrangment(N_AGENT_EACH_TEAM)
    N_TEAM = len(N_AGENT_EACH_TEAM)
    Agent_Property_All = []
    Agent_Property_Array = []

    # chained parameters, will change along with 'N_AGENT_EACH_TEAM'
    AGENT_ID_EACH_TEAM_cv = ChainVar(lambda N_AGENT_EACH_TEAM: usual_id_arrangment(N_AGENT_EACH_TEAM),
                                     chained_with=['N_AGENT_EACH_TEAM'])
    N_TEAM_cv = ChainVar(lambda N_AGENT_EACH_TEAM: len(N_AGENT_EACH_TEAM), chained_with=['N_AGENT_EACH_TEAM'])

    RlAgentEnvironment = "Algorithm.RedModel.CssimAgentEnv.py->CssimAgentEnv"
    # algorithm selection
    Algorithm_Names = ['Algorithm.None->None', ]

    '''
        ## If the length of action array == the number of teams, set ActAsUnity to True
        ## If the length of action array == the number of agents, set ActAsUnity to False
    '''
    ActAsUnity = False

    '''
        ## If the length of reward array == the number of agents, set Enable_Reward_Sharing to False
        ## If the length of reward array == 1, set Enable_Reward_Sharing to True
    '''
    Enable_Reward_Sharing = False

    '''
        ## If the length of obs array == the number of agents, set ObsAsUnity to False
        ## If the length of obs array == the number of teams, set ObsAsUnity to True
    '''
    ObsAsUnity = False

    # <Part 2> Needed by env itself #
    MaxStepsPerEpisode = 100
    render = True  # render with unreal engine client
    js_render = False  # use threejs to run simple render in browser for fast debugging
    # TcpAddr = '127.0.0.1'
    TcpAddr = '192.168.10.15'
    CSSIMPort = 21051

    UnrealLevel = 'CssimUrbanMap'

    # 2025/07: Modify directory structure,New env directory,Modify begin
    # SubTaskSelection = 'Algorithm.RedModel.CssimAgentEnv.py->CssimAgentEnv'
    SubTaskSelection = 'Envs.CssimAgentEnv.py->CssimAgentEnv'
    # 2025/07: Modify end

    # SubTaskSelection = 'CssimAgentEnv'
    SubTaskConfig = get_subtask_conf(SubTaskSelection)
    SubTaskConfig_cv = ChainVar(lambda UnrealLevel: get_subtask_conf(UnrealLevel), chained_with=['SubTaskSelection'])

    UElink2editor = False
    AutoPortOverride = False
    # AutoPortOverride is usually the reverse of UElink2editor
    AutoPortOverride_cv = ChainVar(lambda UElink2editor: (not UElink2editor), chained_with=['UElink2editor'])

    # this is not going to be precise,
    # the precise step time will be floor(StepGameTime/TimeDilation*FrameRate)*TimeDilation/FrameRate
    StepGameTime = 0.5

    CssimServerPath = 'C:/unreal-map/Saved/StagedBuilds/WindowsServer/CSSIMServer.exe'
    UhmapRenderExe = 'C:/unreal-map/Saved/StagedBuilds/WindowsNoEditor'
    TimeDilation = 1.0  # engine calcualtion speed control
    FrameRate = 25.6  # must satisfy: (TimeDilation=1*n, FrameRate=25.6*n)
    FrameRate_cv = ChainVar(lambda TimeDilation: (TimeDilation / 1 * 25.6), chained_with=['TimeDilation'])
    UhmapStartCmd = []
    # <Part 3> Needed by some Algorithm #
    StateProvided = True
    AvailActProvided = False
    EntityOriented = True

    ActionFormat = 'ASCII'  # 'ASCII'/'Multi-Digit'/'Single-Digit'

    n_actions = dictionary_n_actions

    obs_vec_length = get_subtask_conf(SubTaskSelection).obs_vec_length
    obs_vec_length_cv = ChainVar(lambda UnrealLevel: get_subtask_conf(UnrealLevel).obs_vec_length,
                                 chained_with=['SubTaskSelection'])

    obs_n_entity = get_subtask_conf(SubTaskSelection).obs_n_entity
    obs_n_entity_cv = ChainVar(lambda UnrealLevel: get_subtask_conf(UnrealLevel).obs_n_entity,
                               chained_with=['SubTaskSelection'])

    # # ObsBreakBase = 1e4

    UhmapVersion = '3.5'

    CanTurnOff = False

    # Hete agents
    HeteAgents = True

    # 演示demo类别
    DemoType = "Default"


class CSSIMEnvParseHelper:
    def parse_response_ob_info(self, response):
        raise NotImplementedError

    def make_obs(self):
        raise NotImplementedError


class CssimEnv(BaseEnv, CSSIMEnvParseHelper):
    def __init__(self, rank) -> None:
        super().__init__(rank)
        self.id = rank
        self.render = ScenarioConfig.render and (self.id == 0)
        self.n_agents = sum(ScenarioConfig.N_AGENT_EACH_TEAM)
        # 2025/07/04 UE端发送Agent初始化列表，因此不在此处判断数量是否相等
        # assert self.n_agents == len(ScenarioConfig.SubTaskConfig.agent_list), 'agent number defination error'

        self.n_teams = ScenarioConfig.N_TEAM
        self.sim_thread = None
        self.client = None
        # self.observation_space = ?
        # self.action_space = ?
        if ScenarioConfig.StateProvided:
            # self.observation_space['state_shape'] = ?
            pass

        # Restart env, this is very fast, can be a failsafe if there is memory leaking away on UE side
        self.max_simulation_life = 2048

        self.simulation_life = self.max_simulation_life
        # with a lock, we can initialize UE side one by one (not necessary though)

        # wait until thread 0 finish its initialization (to avoid a traffic jam in server memory) 2024/09修改
        # traffic_light = './TEMP/uhmap_thread_0_init_ok_%s'%GlobalConfig.machine_info['ExpUUID'][:8]
        # if rank != 0:
        #     while not os.path.exists(traffic_light): time.sleep(1)

        self.activate_simulation(self.id, find_port=True)

        # thread 0 finish its initialization, 2024/09修改
        # if rank == 0:
        #     with open(traffic_light, mode='w+') as f: f.write(traffic_light)

    def __del__(self):
        self.terminate_simulation()

    def activate_simulation(self, rank, find_port=False):
        # print('thread %d initializing'%rank) # 2024/09 修改
        self.sim_thread = 'activiting'

        # 校验参数
        # print('checking ScenarioConfig args problems ...') # 2024/09 修改
        assert ScenarioConfig.TimeDilation <= 128, "* TimeDilation <= 128 *"
        assert binary_friendly(1 / ScenarioConfig.FrameRate), "* A Butterfly Effect problem *"
        assert binary_friendly(ScenarioConfig.TimeDilation / 256), "* A Butterfly Effect problem *"

        # 获得端口号与启动类型
        self.ue_link_editor = ScenarioConfig.UElink2editor
        self.render = ScenarioConfig.render  # and (rank==0)
        self.hmp_ue_port = ScenarioConfig.CSSIMPort
        self.ip_port = (ScenarioConfig.TcpAddr, self.hmp_ue_port)
        self.FrameRate = 25.6
        self.TimeDilation = 1
        # render和uEditor都是false的话，需要服务器模式启动
        # render是true，uEditor是false的话，固定分配端口，不exe
        # render是false，uEditor是true的话，固定分配端口，不exe
        if self.render:
            # simulation_exe = ScenarioConfig.UhmapRenderExe
            # self.sim_thread = subprocess.Popen([
            #     simulation_exe,
            #     # '-log',
            #     '-TcpPort=%d' % self.hmp_ue_port,  # port for hmp data exchanging
            #     # '-Port=%d' % self.ue_vis_port,  # port for remote visualizing
            #     '-OpenLevel=%s' % ScenarioConfig.UnrealLevel,
            #     '-TimeDilation=%.8f' % ScenarioConfig.TimeDilation,
            #     '-FrameRate=%.8f' % self.FrameRate,
            #     '-IOInterval=%.8f' % ScenarioConfig.StepGameTime,
            #     '-Seed=%d' % int(np.random.rand() * 1e5),  # 如果已经设定了主线程随机数种子，这里随机出来的数字则是确定的
            #     '-DebugMod=False',
            #     # '-LLMCSV',
            #     '-ABSLOG=%s' % os.path.abspath(
            #         './TEMP/cssim/%s/%d.log' % (GlobalConfig.machine_info['ExpUUID'][:8], rank)),
            #     '-Version=%s' % ScenarioConfig.UhmapVersion,
            #     '-LockGameDuringCom=true',
            # ], stdout=subprocess.DEVNULL)
            print(f'Trying to link to CSSIM render with fixed IP {self.ip_port[0]} and port {self.ip_port[1]}...')
        elif self.ue_link_editor:
            print(f'Trying to link to UE editor with fixed IP {self.ip_port[0]} and port {self.ip_port[1]}...')
        else:
            if find_port:
                self.hmp_ue_port, release_port_fn = find_free_port_no_repeat()  # port for hmp data exchanging
                self.ue_vis_port, release_port_fn = find_free_port_no_repeat()  # port for remote visualizing
            simulation_exe = ScenarioConfig.CssimServerPath
            self.ip_port = (ScenarioConfig.TcpAddr, self.hmp_ue_port)
            print(self.ip_port)
            self.sim_thread = subprocess.Popen([
                simulation_exe,
                # '-log',
                '-TcpPort=%d' % self.hmp_ue_port,  # port for hmp data exchanging
                '-Port=%d' % self.ue_vis_port,  # port for remote visualizing
                '-OpenLevel=%s' % ScenarioConfig.UnrealLevel,
                '-TimeDilation=%.8f' % ScenarioConfig.TimeDilation,
                '-FrameRate=%.8f' % ScenarioConfig.FrameRate,
                '-IOInterval=%.8f' % ScenarioConfig.StepGameTime,
                '-Seed=%d' % int(np.random.rand() * 1e5),  # 如果已经设定了主线程随机数种子，这里随机出来的数字则是确定的
                '-DebugMod=False',
                # '-LLMCSV',
                '-ABSLOG=%s' % os.path.abspath(
                    './TEMP/cssim/%s/%d.log' % (GlobalConfig.machine_info['ExpUUID'][:8], rank)),
                '-Version=%s' % ScenarioConfig.UhmapVersion,
                '-LockGameDuringCom=True',
            ], stdout=subprocess.DEVNULL)
            print('To visualize on Windows, run "CSSIMServer.exe -OpenLevel=%s:%d -WINDOWED -TimeDilation=%.8f -FrameRate=%.8f -IOInterval=%.8f -DebugMod=False -LockGameDuringCom=True"' % (get_host_ip(), self.ue_vis_port, ScenarioConfig.TimeDilation, ScenarioConfig.FrameRate,ScenarioConfig.StepGameTime))
            print(f'Trying to link to remote server with path {simulation_exe}...')

        time.sleep(1 + np.abs(self.id) / 100)
        self.client = TcpClientP2PWithCompress(self.ip_port)

        MAX_RETRY = 150
        for i in range(MAX_RETRY):
            try:
                self.client.manual_connect()
                # print('handshake complete %d'%rank)  # 2024/09 修改
                break
            except:
                if i > 25:
                    print(
                        'Thread %d: Trying to connect to unreal engine. Related library not in memory, going to take some minutes. Retry %d ...' % (
                        rank, i))
                elif i > 75:
                    print(
                        'Thread %d: Waiting too long, please reduce parallel threads (num_threads), Retry %d ... | 请减小num_threads运行一次, 让动态库载入内存, 然后恢复num_threads即可' % (
                        rank, i))
                elif i >= MAX_RETRY - 1:
                    assert False, ('cssim connection timeout, please reduce parallel threads (num_threads) !')
                time.sleep(1)
        # now that port is bind, no need to hold them anymore

        # release_port_fn(self.hmp_ue_port)
        if find_port:
            if ScenarioConfig.AutoPortOverride:
                pass
                # release_port_fn(self.hmp_ue_port)
            if not ScenarioConfig.UElink2editor:
                pass
                # release_port_fn(self.ue_vis_port)
        self.t = 0
        # print('thread %d initialize complete'%rank)  # 2024/09 修改

    def terminate_simulation(self):
        if hasattr(self, 'sim_thread') and (self.sim_thread is not None) and (self.client is not None):
            # self.sim_thread.terminate()
            # send terminate command to unreal side
            self.client.send_dgram_to_target(json.dumps({
                'valid': True,
                'DataCmd': 'end_unreal_engine',
                'TimeStepMax': ScenarioConfig.MaxStepsPerEpisode,
                'TimeStep': 0,
                'Actions': None,
            }))
            self.client.close()
            self.sim_thread = None
            self.client = None

    # override reset function
    def reset(self):
        self.simulation_life -= 1
        if self.simulation_life < 0:
            print('restarting simutation')
            self.terminate_simulation()
            self.simulation_life = self.max_simulation_life
            self.activate_simulation(self.id, find_port=False)

    def sleep(self):
        self.simulation_life = -1
        self.terminate_simulation()

    # override step function
    def step(self, act):
        raise NotImplementedError
        # return (ob, RewardForAllTeams,  done, info)  # choose this if Enable_Reward_Sharing
